import logging
import json
from flask import g

from ExtendRegister.model_register import Thesis, Paragraph, GenerateTask
from ExtendRegister.db_register import db
from common import (
    TaskStatus,
    Result,
)
from sqlalchemy import or_
from common.ai import (
    MODEL_GPT35,
    MODEL_GEMINI10,
    MODEL_QIANWEN,
    MODEL_KIMI,
)


from .model.base import BaseAiModel, InvokeResult
from .utils import travelOutlineDict, getModelByName
from .dantic import *
from .model.model_proxy import ModelProxy

from common.task_queue import pubTask, TaskModel, TaskType

LOGGER = logging.getLogger(__name__)


class Service:
    # _init_()方法在Python中为这个类的对象object提供初始化赋值
    def __init__(self) -> None:
        pass

    def getChatReply(self, body: ParamChat) -> Result:
        return Result()

    def exportOutline(self, body: ParamSelect4Content) -> Result:
        allParaTitle = {}

        def _getAllParas(node: dict):
            if node["title"] == "":
                return
            allParaTitle[node["id"]] = node["title"]

        try:

            # 限制一下，同样标题，不能重复提交

            outlineDict = body.outline
            outlineDictWithId = travelOutlineDict(outlineDict, _getAllParas)
            print(outlineDictWithId)
            return Result()
        except Exception as e:
            LOGGER.exception(e)

            return Result().error(str(e))

    def stopThesisGenerate(self, body: ParamThesisId) -> Result:
        try:
            thesisId = body.thesisId
            thesis: Thesis = self.getThesisById(thesisId)
            if thesis is None:
                return Result().error("未查询到论文")

            if thesis.uid != g.userid:
                return Result().error("not found in your list")

            thesis.status = TaskStatus.BEFORE_INIT
            thesis.save()

            cnt = (
                Paragraph.query.filter(Paragraph.thesisId == thesisId)
                .filter(
                    or_(
                        Paragraph.status == TaskStatus.INIT,
                        Paragraph.status == TaskStatus.RUNING,
                    )
                )
                .update({"status": TaskStatus.ERROR})
            )
            db.session.commit()
            LOGGER.info(f"update thesis {thesisId} stop generate")
            return Result()
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def setThesisNeedGenerate(self, thesisId: int) -> Result:
        try:
            thesis: Thesis = Thesis.query.filter(Thesis.id == thesisId).first()

            if thesis.status == TaskStatus.INIT or thesis.status == TaskStatus.RUNING:
                LOGGER.info(f"论文{thesisId}正在生成中，请稍等")
                return Result().error(f"论文正在生成中，请稍等")

            thesis.status = TaskStatus.INIT
            thesis.save()
            LOGGER.info(f"update thesis {thesisId} need update")

            ####
            #### 重要：只有从来没生成过的，和之前生成错误的，可以在这里被在此批量生成
            cnt = (
                Paragraph.query.filter(Paragraph.thesisId == thesisId)
                .filter(
                    or_(
                        Paragraph.status == TaskStatus.BEFORE_INIT,
                        Paragraph.status == TaskStatus.ERROR,
                    )
                )
                .update({"status": TaskStatus.INIT})
            )
            db.session.commit()
            LOGGER.info(f"update Paragraph thesisId:{thesisId} cnt: {cnt}")
            pubTask(TaskModel(thesisId=thesisId, type=TaskType.GEN_THESIS))
            return Result()
        except Exception as e:
            LOGGER.error(e)
            LOGGER.exception(e)
            return Result().error(str(e))

    def generateDigest(self, param: ParamThesisId):
        try:
            t = self.getThesisById(param.thesisId)
            if t is None:
                return Result().error("thesis not found")

            genDigestParam = ParamRegenDigest(
                thesisId=param.thesisId,
                title=t.title,
                outline=t.outline,
                lang=t.lang,
                userId=param.userId,
                isVipUser=param.isVipUser,
            )
            modelProxy = ModelProxy(param.model)
            res: InvokeResult = modelProxy.generateDigest(genDigestParam)
            if not res.isSuccess:
                return Result().error(res.message)

            digestJson = res.data
            # key都定义在userMessage里了，不能轻易变动
            newDigest = digestJson["digest"]
            keywords = digestJson["keywords"]

            if newDigest is None:
                return Result().error("摘要生成失败")

            if newDigest.startswith("摘要："):
                newDigest = newDigest[3:].strip()
            elif newDigest.startswith("摘要:"):
                newDigest = newDigest[3:].strip()
            elif newDigest.startswith("摘要"):
                newDigest = newDigest[2:].strip()

            Thesis.query.filter(Thesis.id == param.thesisId).update(
                {"digest": newDigest, "keywords": keywords}
            )
            db.session.commit()
            return Result()

        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def generateDigestEnglish(self, param: ParamThesisId):
        try:
            t = self.getThesisById(param.thesisId)
            if t is None:
                return Result().error("thesis not found")

            genDigestParam = ParamRegenDigest(
                thesisId=param.thesisId,
                title=t.title,
                outline=t.outline,
                lang=t.lang,
                userId=param.userId,
                isVipUser=param.isVipUser,
            )

            modelProxy = ModelProxy(param.model)
            res: InvokeResult = modelProxy.generateDigestEnglish(genDigestParam)
            if not res.isSuccess:
                return Result().error(res.message)

            digestJson = res.data
            # key都定义在userMessage里了，不能轻易变动
            newDigest = digestJson["digest"]
            keywords = digestJson["keywords"]

            if newDigest is None:
                return Result().error("摘要生成失败")

            if newDigest.startswith("摘要："):
                newDigest = newDigest[3:].strip()
            elif newDigest.startswith("摘要:"):
                newDigest = newDigest[3:].strip()
            elif newDigest.startswith("摘要"):
                newDigest = newDigest[2:].strip()

            Thesis.query.filter(Thesis.id == param.thesisId).update(
                {"digestEn": newDigest, "keywordsEn": keywords}
            )
            db.session.commit()
            return Result()

        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def getParagraphById(self, thesisId: int, paraId: str):
        return (
            Paragraph.query.filter(Paragraph.thesisId == thesisId)
            .filter(Paragraph.paraId == paraId)
            .first()
        )

    def generateSingleParagraph(self, param: ParamGenerateSingleParagraph):
        try:
            t = self.getThesisById(param.thesisId)
            if t is None:
                return Result().error("论文id未找到")

            paragraph: Paragraph = self.getParagraphById(
                param.thesisId, param.paragraphId
            )
            if paragraph is None:
                return Result().error("段落未找到")

            param.paragraphTitle = paragraph.title
            param.title = t.title
            param.lang = t.lang

            mp = ModelProxy(param.model)
            res: InvokeResult = mp.generateParagraph(param)
            if not res.isSuccess:
                return Result().error(res.message)

            newParagraph = res.data
            # 删除句首的关联词
            newParagraph = newParagraph.replace("\n总而言之，", "\n")
            newParagraph = newParagraph.replace("\n综上所述，", "\n")
            newParagraph = newParagraph.replace("\n首先，", "\n")
            newParagraph = newParagraph.replace("\n其次，", "\n")
            newParagraph = newParagraph.replace("\n再次，", "\n")
            newParagraph = newParagraph.replace("\n再者，", "\n")
            newParagraph = newParagraph.replace("\n此外，", "\n")
            newParagraph = newParagraph.replace("\n最后，", "\n")
            newParagraph = newParagraph.replace("\n总之，", "\n")


            Paragraph.query.filter(Paragraph.id == paragraph.id).update(
                {"paragraph": newParagraph, "status": TaskStatus.SUCCESS}
            )
            db.session.commit()
            return Result()

        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def saveThesisInfo(self, thesisId: int, content: str = None):
        try:
            Thesis.query.filter(Thesis.id == thesisId).update({"outline": content})
            db.session.commit()
            return Result()
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

  
    def getThesisById(self, id: int) -> Thesis:
        try:
            return Thesis.query.filter(Thesis.id == id).first()

        except Exception as e:
            LOGGER.exception(e)
            return None

    def select4Content(self, body: ParamSelect4Content) -> Result:

        try:
            if not body.isVipUser:
                if Thesis.query.filter(Thesis.uid == body.userId).count() > 1:
                    return Result().error("抱歉，免费用户最多只能创建一篇论文")
        except Exception as e:
            LOGGER.exception(e)
            return Result().error(f"服务器异常 {str(e)}")

        body.title = body.title.strip()
        if body.title == "":
            return Result().error("请补充论文标题")
        allParaTitle = {}

        def _getAllParas(node: dict):
            if node["title"] == "":
                return
            allParaTitle[node["id"]] = node["title"]

        try:
            userId = g.userid

            # 限制一下，同样标题，不能重复提交

            outlineDict = body.outline
            outlineDictWithId = travelOutlineDict(outlineDict, _getAllParas)

            # 保存提纲，取到论文的id
            t = Thesis(
                uid=userId,
                title=body.title,
                outline=outlineDictWithId,
                lang=body.lang,
                level=body.level,
                length=body.length,
            )

            db.session.add(t)
            db.session.commit()

            # 逐个保存段落，进入Paragraph表
            for id, title in allParaTitle.items():
                db.session.add(
                    Paragraph(uid=userId, thesisId=t.id, paraId=id, title=title)
                )
            db.session.commit()

            return Result().setData({"thesisId": t.id})
        except Exception as e:
            db.session.rollback()
            LOGGER.exception(e)
            return Result().error(str(e))

    # 传入的两个参数，self参数用于获得AI模型，param: ParamOutline参数表示，期待传入一个ParamOutline的类的对象作为参数param
    # 前端需要保证，前端通过axios提交Http Request传过来的参数和后端的ParamOutline类定义的格式一致。
    # 这个代码实际上使用了FastAPI,服务器会自动解析HTTP Request类，拿到服务器端处理方法所需要的参数
    # 也就是ParamOutline格式的类param，如果没有这个类似的参数，服务器会直接拒绝服务，处理的方法getOutline也不会被调用

    def getOutline(self, param: ParamOutline) -> Result:

        try:

            param.title = param.title.strip()
            if param.title == "":
                return Result().error("请填写论文标题")

            # 非VIP 章节只能到6，vip可以到12。其实页面上写的是9
            if param.isVipUser:
                param.paragraphCount = (
                    12
                    if param.paragraphCount > 12 or param.paragraphCount < 1
                    else param.paragraphCount
                )
            else:
                param.paragraphCount = (
                    6
                    if param.paragraphCount > 6 or param.paragraphCount < 1
                    else param.paragraphCount
                )
            LOGGER.info(f"request ${param.model} param: {param}")
            m = ModelProxy(param.model)
            res: InvokeResult = m.getOutline(param)
            if res.isSuccess:
                return Result().setData(res.data)
            return Result().error(res.message)

        except Exception as e:
            LOGGER.exception(e)
            return Result().error(str(e))

    def getTitle(self, param: ParamTitle) -> Result:

        try:
            invokeResult: InvokeResult = getModelByName(param.model).getTitle(param)
            if invokeResult.isSuccess:
                return Result().setData(json.loads(invokeResult.data))

            return Result().error(invokeResult.message)

        except Exception as e:

            LOGGER.exception(e)
            return Result().error(str(e))
